import torch
import numpy as np
from local_occ_grid_map import LocalMap
from torch.utils.data import Dataset, DataLoader
import os
import imageio as iio


class OGMImageDataset(Dataset):
    def __init__(self, name, base_path, split='train'):
        self.name = name
        self.split = split
        scan_base_path = os.path.join(base_path, split, 'scans')
        pos_base_path = os.path.join(base_path, split, 'positions')
        vel_base_path = os.path.join(base_path, split, 'velocities')
        scan_meta = os.path.join(scan_base_path, split+'.txt')
        pos_meta = os.path.join(pos_base_path, split+'.txt')
        vel_meta = os.path.join(vel_base_path, split+'.txt')
        self.scan_files, self.pos_files, self.vel_files = [], [], []
        with open(scan_meta, 'r') as f:
            for line in f.read().split('\n'):
                if '.npy' in line:
                    self.scan_files.append(os.path.join(scan_base_path, line))
        with open(pos_meta, 'r') as f:
            for line in f.read().split('\n'):
                if '.npy' in line:
                    self.pos_files.append(os.path.join(pos_base_path, line))
        with open(vel_meta, 'r') as f:
            for line in f.read().split('\n'):
                if '.npy' in line:
                    self.vel_files.append(os.path.join(vel_base_path, line))
        self.length = len(self.scan_files)
        
    def __len__(self):
        return self.length
    
    def __getitem__(self, index):
        scan = np.load(self.scan_files[index])
        map_name = self.scan_files[index].split('.')[0].split('/')[-1] + '.png'
        mask_gridMap = LocalMap(X_lim = [0, 6.4], 
                        Y_lim = [-3.2, 3.2], 
                        resolution = 0.1, 
                        p = 0.5,
                        size=[1, 1],
                        device = torch.device('cpu'))
        # robot positions:
        x_odom = torch.zeros(1, 1)
        y_odom = torch.zeros(1, 1)
        theta_odom = torch.zeros(1, 1)
        # Lidar measurements:
        distances = torch.from_numpy(scan[None, None, ...])
        # the angles of lidar scan: -135 ~ 135 degree
        angles = torch.linspace(-(135*np.pi/180), 135*np.pi/180, distances.shape[-1])
        # Lidar measurements in X-Y plane: transform to the predicted robot reference frame
        distances_x, distances_y = mask_gridMap.lidar_scan_xy(distances, angles, x_odom, y_odom, theta_odom)
        # discretize to binary maps:
        mask_binary_maps = mask_gridMap.discretize(distances_x, distances_y)
        data = {
            'map': mask_binary_maps,
            'name': map_name
        }
        return data


def create_maps(dataset, num_workers=8):
    dest_path = os.path.join(dataset.name + '-maps', dataset.split)
    os.makedirs(dest_path, exist_ok=True)
    dloader = DataLoader(dataset, batch_size=1, shuffle=False, drop_last=False, num_workers=num_workers)
    for batch in dloader:
        map_path = os.path.join(dest_path, batch['name'][0])
        map_tensor = batch['map']
        map_todump = map_tensor.squeeze().numpy().astype(np.uint8) * 255
        iio.imwrite(map_path, map_todump)
        print(f'Save map to {map_path}')

if __name__ == "__main__":
    dset_train = OGMImageDataset('OGM-Turtlebot2' ,'OGM-datasets/OGM-Turtlebot2')
    dset_val = OGMImageDataset('OGM-Turtlebot2' ,'OGM-datasets/OGM-Turtlebot2', 'val')
    dset_test = OGMImageDataset('OGM-Turtlebot2' ,'OGM-datasets/OGM-Turtlebot2', 'test')
    create_maps(dset_train)
    create_maps(dset_val)
    create_maps(dset_test)